import os
from collections import defaultdict
from pke_zh import KeyBert

# 初始化 KeyBert
keybert_extractor = KeyBert()

# 倒排索引字典
inverted_index = defaultdict(list)

# 指定文档所在的相对路径
document_directory = '../data/corpus'

# 获取目录中的所有文档路径
documents = [os.path.join(document_directory, filename) for filename in os.listdir(document_directory) if filename.endswith('.txt')]

print("Current working directory:", os.getcwd())
print("Document paths:", documents)

# 遍历文档，提取关键词并更新倒排索引
for doc_path in documents:
    with open(doc_path, 'r', encoding='utf-8') as file:
        content = file.read()
        doc_id = os.path.basename(doc_path)
        keywords = keybert_extractor.extract(content, n_best=10, use_maxsum=True, alpha=0.5)
        for keyword, score in keywords:
            inverted_index[keyword].append((doc_id, score))

# 打印或保存倒排索引
for keyword, docs in inverted_index.items():
    print(f"Keyword: {keyword}, Documents: {docs}")
